import heapq
import time
import numpy as np


class Node:
    def __init__(self, state, parent=None, g=0, h=0):
        self.state = state  # 3x3矩阵表示状态
        self.parent = parent  # 父节点
        self.g = g  # 实际代价
        self.h = h  # 启发式函数值
        self.f = g + h  # 估价函数

    def __lt__(self, other):
        return self.f < other.f  # 优先队列按f值排序


def manhattan_distance(state, target):
    """计算曼哈顿距离"""
    distance = 0
    for i in range(3):
        for j in range(3):
            if state[i][j] == 0:
                continue
            value = state[i][j]
            target_pos = np.argwhere(target == value)[0]
            distance += abs(i - target_pos[0]) + abs(j - target_pos[1])
    return distance


def misplaced_tiles(state, target):
    """计算错位数"""
    return np.sum(state != target) - (state == 0).sum()  # 排除空格


def get_children(node, target, heuristic):
    """生成子节点"""
    children = []
    zero_pos = np.argwhere(node.state == 0)[0]
    x, y = zero_pos

    # 四个可能的移动方向
    moves = [(-1, 0), (1, 0), (0, -1), (0, 1)]
    for dx, dy in moves:
        nx, ny = x + dx, y + dy
        if 0 <= nx < 3 and 0 <= ny < 3:
            new_state = np.copy(node.state)
            new_state[x][y], new_state[nx][ny] = new_state[nx][ny], new_state[x][y]

            # 计算启发式函数值
            if heuristic == "manhattan":
                h = manhattan_distance(new_state, target)
            else:
                h = misplaced_tiles(new_state, target)

            child = Node(new_state, node, node.g + 1, h)
            children.append(child)
    return children


def a_star(initial, target, heuristic="manhattan"):
    """A*算法实现"""
    open_list = []
    close_set = set()

    initial_node = Node(np.array(initial), None, 0, 0)
    target = np.array(target)

    # 初始化启发式函数值
    if heuristic == "manhattan":
        initial_node.h = manhattan_distance(initial_node.state, target)
    else:
        initial_node.h = misplaced_tiles(initial_node.state, target)
    initial_node.f = initial_node.g + initial_node.h

    heapq.heappush(open_list, initial_node)
    expanded_nodes = 0

    while open_list:
        current_node = heapq.heappop(open_list)
        expanded_nodes += 1

        # 到达目标状态
        if np.array_equal(current_node.state, target):
            path = []
            while current_node:
                path.append(current_node.state)
                current_node = current_node.parent
            return path[::-1], len(path) - 1, expanded_nodes  # 反转路径

        # 将当前状态加入Close表
        state_key = tuple(map(tuple, current_node.state))
        if state_key in close_set:
            continue
        close_set.add(state_key)

        # 生成子节点
        children = get_children(current_node, target, heuristic)
        for child in children:
            child_key = tuple(map(tuple, child.state))
            if child_key not in close_set:
                heapq.heappush(open_list, child)

    return None, 0, expanded_nodes  # 无解


# 示例输入
initial_state = [[2, 8, 3], [1, 6, 4], [7, 0, 5]]
target_state = [[1, 2, 3], [8, 0, 4], [7, 6, 5]]

# 运行A*算法（曼哈顿距离）
start_time = time.time()
path, steps, expanded = a_star(initial_state, target_state, heuristic="manhattan")
end_time = time.time()

# 输出结果
if path:
    print("最优路径（部分步骤）：")
    for i, state in enumerate(path):
        print(f"步骤{i}:")
        print(np.array(state))
    print(f"总步数: {steps}")
    print(f"扩展节点数: {expanded}")
    print(f"运行时间: {end_time - start_time:.4f}秒")
else:
    print("无解")